package com.lord.otter.component.utils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 * 树工具类
 *
 * @author lord
 * 2025/8/27 13:50
 */
public class TreeKit {

    /**
     * 平铺列表转树
     *
     * @param flatList            数据源
     * @param idGetter            获取主键
     * @param parentIdGetter      获取父节点ID的函数
     * @param childrenGetter      获取子列表的函数
     * @param childrenInitializer 初始化空子列表的函数
     * @param childAdder          添加子节点的函数
     * @author lord
     * 2025/8/27 16:36
     */
    public static <T, K> List<T> build(
            List<T> flatList,
            Function<T, K> idGetter,
            Function<T, K> parentIdGetter,
            Function<T, List<T>> childrenGetter,
            Consumer<T> childrenInitializer,
            Consumer<NodeRelation<T>> childAdder
    ) {
        // 1. 创建ID映射表
        Map<K, T> idMap = new HashMap<>(flatList.size());
        for (T node : flatList) {
            K id = idGetter.apply(node);
            if (id == null) {
                throw new IllegalArgumentException("Node ID cannot be null: " + node);
            }
            idMap.put(id, node);
        }
        // 2. 构建树结构
        List<T> roots = new ArrayList<>();
        for (T node : flatList) {
            K parentId = parentIdGetter.apply(node);

            if (parentId == null) {
                // 父ID为null，视为根节点
                roots.add(node);
            } else {
                T parent = idMap.get(parentId);
                if (parent != null) {
                    // 懒初始化子列表
                    if (childrenGetter.apply(parent) == null) {
                        childrenInitializer.accept(parent);
                    }
                    // 添加父子关系
                    childAdder.accept(new NodeRelation<>(parent, node));
                } else {
                    // 父节点不存在，视为根节点
                    roots.add(node);
                }
            }
        }
        return roots;
    }

    /**
     * 用于封装父子关系的值对象
     */
    public static class NodeRelation<T> {
        public final T parent;
        public final T child;

        public NodeRelation(T parent, T child) {
            this.parent = parent;
            this.child = child;
        }
    }
}
